/*
 * Copyright 2021-2022 Open Kunlun Technology <https://www.openkunlun.io>
 */

package io.openkunlun.scaladsl.context

import akka.actor.{ ExtendedActorSystem, Extension, ExtensionId, ExtensionIdProvider }
import akka.event.Logging
import io.openkunlun.scaladsl.context.Context.ContextTag
import kamon.Kamon
import kamon.context.{ Context => KamonContext }
import kamon.tag.{ Tag, TagSet }

import scala.util.Try

/**
 * @author ericxin.
 */

private object KamonContextStorage extends ExtensionId[KamonContextStorage] with ExtensionIdProvider {
  val Name = "kamon"
  override def lookup: ExtensionId[KamonContextStorage] = KamonContextStorage
  override def createExtension(system: ExtendedActorSystem) = new KamonContextStorage(system)
}
private class KamonContextStorage(system: ExtendedActorSystem) extends Extension with ContextStorage {

  private val log = Logging(system, getClass)

  var kamonConfig = Kamon.config()
  if (system.settings.config.getBoolean("io.openkunlun.context.instrumentation.akka-http")) {
    kamonConfig = kamonConfig.withFallback(system.settings.config.getConfig(s"io.openkunlun.context.${KamonContextStorage.Name}.akka-http"))
  }
  if (system.settings.config.getBoolean("io.openkunlun.context.instrumentation.play")) {
    kamonConfig = kamonConfig.withFallback(system.settings.config.getConfig(s"io.openkunlun.context.${KamonContextStorage.Name}.play"))
  }

  override def init(): Unit = {
    log.info("Init kamon context storage.")
    Kamon.init(kamonConfig)
  }

  override def store[T](ctx: Context)(f: => T): T = {
    val builder = TagSet.builder()
    ctx.tags.map(it => it.key -> it.value)
      .foreach(it => builder.add(it._1, it._2))
    Kamon.runWithContext(KamonContext.of(builder.build()))(f)
  }

  override def get(): Context = {
    val tags = Kamon.currentContext().tags
      .iterator()
      .map(it => it.key -> Try(Tag.unwrapValue(it).toString).getOrElse(""))
      .map(it => ContextTag(it._1, it._2))
      .toSeq
    new Context(tags)
  }
}
